{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Copyright (c) Microsoft Corporation. All rights reserved.\n",
    "\n",
    "Licensed under the MIT License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Question Answering on the SQuAD Dataset using Transformers Models\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Before You Start\n",
    "\n",
    "The running time shown in this notebook is on a Standard_NC24rs_v3 Azure Data Science Virtual Machine with 4 NVIDIA Tesla V100 GPUs. \n",
    "> **Tip**: If you want to run through the notebook quickly, you can set the **`QUICK_RUN`** flag in the cell below to **`True`** to run the notebook on a small subset of the data and a smaller number of epochs. \n",
    "\n",
    "The table below provides some reference running time of BERT on different machine configurations.  \n",
    "\n",
    "|QUICK_RUN|Machine Configurations|Running time|\n",
    "|:---------|:----------------------|:------------|\n",
    "|True|4 **CPU**s, 14GB memory| ~ 10 minutes |\n",
    "|True|1 NVIDIA Tesla K80 GPUs, 12GB GPU memory| ~ 3 minutes |\n",
    "|False|4 NVIDIA Tesla K80 GPUs, 48GB GPU memory| ~ 18 hours |\n",
    "|False|4 NVIDIA Tesla V100 GPUs, 64GB GPU memory, without RDMA (NC24s)| ~ 7 hours|\n",
    "|False|4 NVIDIA Tesla V100 GPUs, 64GB GPU memory, with RDMA (NC24**r**s)| ~ 4 hours|\n",
    "\n",
    "If you run into CUDA out-of-memory error, try reducing the `PER_GPU_BATCH_SIZE` and increasing the `GRADIENT_ACCUMULATION_STEPS`. As long as `PER_GPU_BATCH_SIZE` * `GRADIENT_ACCUMULATION_STEPS` doesn't change, the effective **per gpu** batch size is the same as larger `PER_GPU_BATCH_SIZE` and smaller `GRADIENT_ACCUMULATION_STEPS`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Set QUICK_RUN = True to run the notebook on a small subset of data and a smaller number of epochs.\n",
    "QUICK_RUN = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "This notebook demonstrates how to fine tune [pre-trained transformers models](https://github.com/huggingface/transformers) for extractive question answering task. Utility functions and classes in the NLP Best Practices repo are used to facilitate data preprocessing, model training, model scoring, result postprocessing, and model evaluation. \n",
    "\n",
    "The following models are currently supported:\n",
    "* BERT: [Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805)\n",
    "* XLNet: [Generalized Autoregressive Pretraining for Language Understanding](https://arxiv.org/pdf/1906.08237.pdf)\n",
    "* DistilBert: [A small, fast, cheap and light Transformer model based on Bert architecture](https://medium.com/huggingface/distilbert-8cf3380435b5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "\n",
    "import scrapbook as sb\n",
    "import torch\n",
    "\n",
    "from utils_nlp.common.pytorch_utils import dataloader_from_dataset\n",
    "from utils_nlp.common.timer import Timer\n",
    "from utils_nlp.dataset.squad import load_pandas_df\n",
    "from utils_nlp.eval.question_answering import evaluate_qa\n",
    "from utils_nlp.models.transformers.datasets import QADataset\n",
    "from utils_nlp.models.transformers.question_answering import (\n",
    "    AnswerExtractor,\n",
    "    QAProcessor,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configurations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To get all the transformer models supporting question answering, call `AnswerExtractor.list_supported_models()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['bert-base-uncased',\n",
       " 'bert-large-uncased',\n",
       " 'bert-base-cased',\n",
       " 'bert-large-cased',\n",
       " 'bert-base-multilingual-uncased',\n",
       " 'bert-base-multilingual-cased',\n",
       " 'bert-base-chinese',\n",
       " 'bert-base-german-cased',\n",
       " 'bert-large-uncased-whole-word-masking',\n",
       " 'bert-large-cased-whole-word-masking',\n",
       " 'bert-large-uncased-whole-word-masking-finetuned-squad',\n",
       " 'bert-large-cased-whole-word-masking-finetuned-squad',\n",
       " 'bert-base-cased-finetuned-mrpc',\n",
       " 'xlnet-base-cased',\n",
       " 'xlnet-large-cased',\n",
       " 'distilbert-base-uncased',\n",
       " 'distilbert-base-uncased-distilled-squad']"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "AnswerExtractor.list_supported_models()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": [
     "parameters"
    ]
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Max sequence length: 384\n",
      "Document stride: 128\n",
      "Per gpu batch size: 4\n"
     ]
    }
   ],
   "source": [
    "MODEL_NAME = \"bert-large-cased-whole-word-masking\"\n",
    "DO_LOWER_CASE = False\n",
    "\n",
    "# MODEL_NAME = \"xlnet-large-cased\"\n",
    "# DO_LOWER_CASE = False\n",
    "\n",
    "# MODEL_NAME = \"distilbert-base-uncased\"\n",
    "# DO_LOWER_CASE = True\n",
    "\n",
    "TRAIN_DATA_USED_PERCENT = 1\n",
    "DEV_DATA_USED_PERCENT = 1\n",
    "NUM_EPOCHS = 2\n",
    "\n",
    "MAX_SEQ_LENGTH = 384\n",
    "DOC_STRIDE = 128\n",
    "PER_GPU_BATCH_SIZE = 4\n",
    "GRADIENT_ACCUMULATION_STEPS = 1\n",
    "NUM_GPUS = torch.cuda.device_count()\n",
    "\n",
    "if QUICK_RUN:\n",
    "    TRAIN_DATA_USED_PERCENT = 0.001\n",
    "    DEV_DATA_USED_PERCENT = 0.01\n",
    "    NUM_EPOCHS = 1\n",
    "    \n",
    "    MAX_SEQ_LENGTH = 128\n",
    "    DOC_STRIDE = 64\n",
    "    PER_GPU_BATCH_SIZE = 1\n",
    "\n",
    "print(\"Max sequence length: {}\".format(MAX_SEQ_LENGTH))\n",
    "print(\"Document stride: {}\".format(DOC_STRIDE))\n",
    "print(\"Per gpu batch size: {}\".format(PER_GPU_BATCH_SIZE))\n",
    "\n",
    "RANDOM_SEED = 42\n",
    "SQUAD_VERSION = \"v1.1\" \n",
    "CACHE_DIR = \"./temp\"\n",
    "\n",
    "MAX_QUESTION_LENGTH = 64\n",
    "LEARNING_RATE = 3e-5\n",
    "\n",
    "DOC_TEXT_COL = \"doc_text\"\n",
    "QUESTION_TEXT_COL = \"question_text\"\n",
    "ANSWER_START_COL = \"answer_start\"\n",
    "ANSWER_TEXT_COL = \"answer_text\"\n",
    "QA_ID_COL = \"qa_id\"\n",
    "IS_IMPOSSIBLE_COL = \"is_impossible\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### The SQuAD Dataset\n",
    "Stanford Question Answering Dataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable. [\\[1, 2\\]](#References)\n",
    "\n",
    "<img src=\"https://nlpbp.blob.core.windows.net/images/squad.png\">\n",
    "\n",
    "There has been two versions of SQuAD datasets. SQuAD 1.1 contains 100,000+ question-answer pairs on 500+ articles. SQuAD 2.0 adds 50,000 new, unanswerable questions written adversarially by crowdworkers to look similar to answerable ones. These datasets are available at [https://rajpurkar.github.io/SQuAD-explorer/](https://rajpurkar.github.io/SQuAD-explorer/). Each dataset comes with a training dataset and a development dataset. \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The utility function `load_pandas_df` downloads the dataset specified by `squad_version` and `file_split` to `local_cache_path` if it doesn't exist already."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 7.82k/7.82k [00:00<00:00, 20.6kKB/s]\n",
      "100%|██████████| 1.02k/1.02k [00:00<00:00, 19.9kKB/s]\n"
     ]
    }
   ],
   "source": [
    "train_df = load_pandas_df(local_cache_path=CACHE_DIR, squad_version=SQUAD_VERSION, file_split=\"train\")\n",
    "dev_df = load_pandas_df(local_cache_path=CACHE_DIR, squad_version=SQUAD_VERSION, file_split=\"dev\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>doc_text</th>\n",
       "      <th>question_text</th>\n",
       "      <th>answer_start</th>\n",
       "      <th>answer_text</th>\n",
       "      <th>qa_id</th>\n",
       "      <th>is_impossible</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>To whom did the Virgin Mary allegedly appear i...</td>\n",
       "      <td>515</td>\n",
       "      <td>Saint Bernadette Soubirous</td>\n",
       "      <td>5733be284776f41900661182</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What is in front of the Notre Dame Main Building?</td>\n",
       "      <td>188</td>\n",
       "      <td>a copper statue of Christ</td>\n",
       "      <td>5733be284776f4190066117f</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>The Basilica of the Sacred heart at Notre Dame...</td>\n",
       "      <td>279</td>\n",
       "      <td>the Main Building</td>\n",
       "      <td>5733be284776f41900661180</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What is the Grotto at Notre Dame?</td>\n",
       "      <td>381</td>\n",
       "      <td>a Marian place of prayer and reflection</td>\n",
       "      <td>5733be284776f41900661181</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>Architecturally, the school has a Catholic cha...</td>\n",
       "      <td>What sits on top of the Main Building at Notre...</td>\n",
       "      <td>92</td>\n",
       "      <td>a golden statue of the Virgin Mary</td>\n",
       "      <td>5733be284776f4190066117e</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            doc_text  \\\n",
       "0  Architecturally, the school has a Catholic cha...   \n",
       "1  Architecturally, the school has a Catholic cha...   \n",
       "2  Architecturally, the school has a Catholic cha...   \n",
       "3  Architecturally, the school has a Catholic cha...   \n",
       "4  Architecturally, the school has a Catholic cha...   \n",
       "\n",
       "                                       question_text  answer_start  \\\n",
       "0  To whom did the Virgin Mary allegedly appear i...           515   \n",
       "1  What is in front of the Notre Dame Main Building?           188   \n",
       "2  The Basilica of the Sacred heart at Notre Dame...           279   \n",
       "3                  What is the Grotto at Notre Dame?           381   \n",
       "4  What sits on top of the Main Building at Notre...            92   \n",
       "\n",
       "                               answer_text                     qa_id  \\\n",
       "0               Saint Bernadette Soubirous  5733be284776f41900661182   \n",
       "1                a copper statue of Christ  5733be284776f4190066117f   \n",
       "2                        the Main Building  5733be284776f41900661180   \n",
       "3  a Marian place of prayer and reflection  5733be284776f41900661181   \n",
       "4       a golden statue of the Virgin Mary  5733be284776f4190066117e   \n",
       "\n",
       "   is_impossible  \n",
       "0          False  \n",
       "1          False  \n",
       "2          False  \n",
       "3          False  \n",
       "4          False  "
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>doc_text</th>\n",
       "      <th>question_text</th>\n",
       "      <th>answer_start</th>\n",
       "      <th>answer_text</th>\n",
       "      <th>qa_id</th>\n",
       "      <th>is_impossible</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>0</td>\n",
       "      <td>Super Bowl 50 was an American football game to...</td>\n",
       "      <td>Which NFL team represented the AFC at Super Bo...</td>\n",
       "      <td>[177, 177, 177]</td>\n",
       "      <td>[Denver Broncos, Denver Broncos, Denver Broncos]</td>\n",
       "      <td>56be4db0acb8001400a502ec</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>Super Bowl 50 was an American football game to...</td>\n",
       "      <td>Which NFL team represented the NFC at Super Bo...</td>\n",
       "      <td>[249, 249, 249]</td>\n",
       "      <td>[Carolina Panthers, Carolina Panthers, Carolin...</td>\n",
       "      <td>56be4db0acb8001400a502ed</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>Super Bowl 50 was an American football game to...</td>\n",
       "      <td>Where did Super Bowl 50 take place?</td>\n",
       "      <td>[403, 355, 355]</td>\n",
       "      <td>[Santa Clara, California, Levi's Stadium, Levi...</td>\n",
       "      <td>56be4db0acb8001400a502ee</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>Super Bowl 50 was an American football game to...</td>\n",
       "      <td>Which NFL team won Super Bowl 50?</td>\n",
       "      <td>[177, 177, 177]</td>\n",
       "      <td>[Denver Broncos, Denver Broncos, Denver Broncos]</td>\n",
       "      <td>56be4db0acb8001400a502ef</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>Super Bowl 50 was an American football game to...</td>\n",
       "      <td>What color was used to emphasize the 50th anni...</td>\n",
       "      <td>[488, 488, 521]</td>\n",
       "      <td>[gold, gold, gold]</td>\n",
       "      <td>56be4db0acb8001400a502f0</td>\n",
       "      <td>False</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                            doc_text  \\\n",
       "0  Super Bowl 50 was an American football game to...   \n",
       "1  Super Bowl 50 was an American football game to...   \n",
       "2  Super Bowl 50 was an American football game to...   \n",
       "3  Super Bowl 50 was an American football game to...   \n",
       "4  Super Bowl 50 was an American football game to...   \n",
       "\n",
       "                                       question_text     answer_start  \\\n",
       "0  Which NFL team represented the AFC at Super Bo...  [177, 177, 177]   \n",
       "1  Which NFL team represented the NFC at Super Bo...  [249, 249, 249]   \n",
       "2                Where did Super Bowl 50 take place?  [403, 355, 355]   \n",
       "3                  Which NFL team won Super Bowl 50?  [177, 177, 177]   \n",
       "4  What color was used to emphasize the 50th anni...  [488, 488, 521]   \n",
       "\n",
       "                                         answer_text  \\\n",
       "0   [Denver Broncos, Denver Broncos, Denver Broncos]   \n",
       "1  [Carolina Panthers, Carolina Panthers, Carolin...   \n",
       "2  [Santa Clara, California, Levi's Stadium, Levi...   \n",
       "3   [Denver Broncos, Denver Broncos, Denver Broncos]   \n",
       "4                                 [gold, gold, gold]   \n",
       "\n",
       "                      qa_id  is_impossible  \n",
       "0  56be4db0acb8001400a502ec          False  \n",
       "1  56be4db0acb8001400a502ed          False  \n",
       "2  56be4db0acb8001400a502ee          False  \n",
       "3  56be4db0acb8001400a502ef          False  \n",
       "4  56be4db0acb8001400a502f0          False  "
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dev_df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_df = train_df.sample(frac=TRAIN_DATA_USED_PERCENT).reset_index(drop=True)\n",
    "dev_df = dev_df.sample(frac=DEV_DATA_USED_PERCENT).reset_index(drop=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`QADataset` is a standard question answering dataset for downstream processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = QADataset(\n",
    "    df=train_df,\n",
    "    doc_text_col=DOC_TEXT_COL,\n",
    "    question_text_col=QUESTION_TEXT_COL,\n",
    "    qa_id_col=QA_ID_COL,\n",
    "    is_impossible_col=IS_IMPOSSIBLE_COL,\n",
    "    answer_start_col=ANSWER_START_COL,\n",
    "    answer_text_col=ANSWER_TEXT_COL\n",
    ")\n",
    "dev_dataset = QADataset(\n",
    "    df=dev_df,\n",
    "    doc_text_col=DOC_TEXT_COL,\n",
    "    question_text_col=QUESTION_TEXT_COL,\n",
    "    qa_id_col=QA_ID_COL,\n",
    "    is_impossible_col=IS_IMPOSSIBLE_COL,\n",
    "    answer_start_col=ANSWER_START_COL,\n",
    "    answer_text_col=ANSWER_TEXT_COL\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Tokenize and Preprocess Data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `QAProcessor.preprocess` tokenizes the input paragraph, question, and answer texts, and converts them into the format required by pre-trained transformer models, involving the following steps:\n",
    "* Tokenization.\n",
    "* Convert character-based answer span indices to token-based indices.\n",
    "* Truncate the question token list if it's longer than `max_question_length`.\n",
    "* Split the paragraph into multiple segments if it's longer than `max_seq_length` - `max_question_length` - 3. (The \"-3\" is for the special [CLS] token and two [SEP] tokens.)\n",
    "* Add the special tokens [CLS] and [SEP].\n",
    "* Pad the concatenated token sequence to `max_seq_length` if it's shorter.\n",
    "* Convert the tokens into token indices corresponding to the tokenizer's vocabulary.\n",
    "\n",
    "`QAProcessor.preprocess` returns a Pytorch DataSet. By default, it saves `cached_examples_train/test.jsonl` and `cached_features_train/test.jsonl` to `./cached_qa_features`. These files are required by postprocessing the predicted answer start and end indices to get the final answer text. You can change the default file directory by specifying `feature_cache_dir`. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 213450/213450 [00:00<00:00, 2918674.41B/s]\n"
     ]
    }
   ],
   "source": [
    "qa_processor = QAProcessor(model_name=MODEL_NAME, to_lower=DO_LOWER_CASE)\n",
    "train_dataset = qa_processor.preprocess(\n",
    "    train_dataset,\n",
    "    is_training=True,\n",
    "    max_question_length=MAX_QUESTION_LENGTH,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    doc_stride=DOC_STRIDE,\n",
    ")\n",
    "\n",
    "# we keep a copy of the oroginal dev_dataset as it is needed for evaluation\n",
    "dev_dataset_processed = qa_processor.preprocess(\n",
    "    dev_dataset,\n",
    "    is_training=False,\n",
    "    max_question_length=MAX_QUESTION_LENGTH,\n",
    "    max_seq_length=MAX_SEQ_LENGTH,\n",
    "    doc_stride=DOC_STRIDE,\n",
    ")\n",
    "\n",
    "train_dataloader = dataloader_from_dataset(\n",
    "    train_dataset, batch_size=PER_GPU_BATCH_SIZE, num_gpus=NUM_GPUS, shuffle=True\n",
    ")\n",
    "dev_dataloader = dataloader_from_dataset(\n",
    "    dev_dataset_processed, batch_size=PER_GPU_BATCH_SIZE, num_gpus=NUM_GPUS, shuffle=False\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fine-tune AnswerExtractor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_extractor = AnswerExtractor(model_name=MODEL_NAME, cache_dir=CACHE_DIR)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with Timer() as t:\n",
    "    qa_extractor.fit(train_dataloader,\n",
    "                     num_epochs=NUM_EPOCHS,\n",
    "                     learning_rate=LEARNING_RATE,\n",
    "                     gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,\n",
    "                     seed=RANDOM_SEED,\n",
    "                     cache_model=True)\n",
    "print(\"Training time : {:.3f} hrs\".format(t.interval / 3600)) "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Predict\n",
    "Note that the `AnswerExtractor.predict` only outputs the probabilities of each token being the start and end of the answer span. `postprocess_bert_answer` and  `postprocess_xlnet_answer` are two helper functions for postprocessing these probabilities and generating the final answers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 661/661 [04:42<00:00,  2.64it/s]\n"
     ]
    }
   ],
   "source": [
    "qa_results = qa_extractor.predict(dev_dataloader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Postprocess and Generate the Final Answers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "final_answers, answer_probs, nbest_answers = qa_processor.postprocess(\n",
    "    qa_results,\n",
    "    examples_file=\"./cached_qa_features/cached_examples_test.jsonl\",\n",
    "    features_file=\"./cached_qa_features/cached_features_test.jsonl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Paragraph:\n",
      "In August 1999, ABC premiered a special series event, Who Wants to Be a Millionaire, a game show based on the British program of the same title. Hosted throughout its ABC tenure by Regis Philbin, the program became a major ratings success throughout its initial summer run, which led ABC to renew Millionaire as a regular series, returning on January 18, 2000. At its peak, the program aired as much as six nights a week. Buoyed by Millionaire, during the 1999–2000 season, ABC became the first network to move from third to first place in the ratings during a single television season. Millionaire ended its run on the network's primetime lineup after three years in 2002, with Buena Vista Television relaunching the show as a syndicated program (under that incarnation's original host Meredith Vieira) in September of that year.\n",
      "\n",
      "Question:\n",
      "Who originally hosted Who Wants to Be a Millionaire for ABC?\n",
      "\n",
      "Ground truth answers:\n",
      "['Regis Philbin', 'Regis Philbin', 'Regis Philbin']\n",
      "\n",
      "Predicted answer:\n",
      "Regis Philbin\n",
      "\n",
      "Top N best answers\n",
      "[OrderedDict([('text', 'Regis Philbin'), ('probability', 0.9916362000840248), ('start_logit', 7.9423394203186035), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Hosted throughout its ABC tenure by Regis Philbin'), ('probability', 0.004141487486401103), ('start_logit', 2.464038133621216), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Regis Philbin,'), ('probability', 0.0035518662225726477), ('start_logit', 7.9423394203186035), ('end_logit', 1.9868273735046387)]), OrderedDict([('text', 'Philbin'), ('probability', 0.0003250549061321865), ('start_logit', -0.08077805489301682), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Regis'), ('probability', 0.00019582582917568986), ('start_logit', 7.9423394203186035), ('end_logit', -0.9111754298210144)]), OrderedDict([('text', 'Regis Phil'), ('probability', 4.149332683962915e-05), ('start_logit', 7.9423394203186035), ('end_logit', -2.4628684520721436)]), OrderedDict([('text', 'by Regis Philbin'), ('probability', 3.332539919374182e-05), ('start_logit', -2.358452320098877), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'bin'), ('probability', 2.8618519913389162e-05), ('start_logit', -2.5107181072235107), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Hosted throughout its ABC tenure by Regis Philbin,'), ('probability', 1.483407878101762e-05), ('start_logit', 2.464038133621216), ('end_logit', 1.9868273735046387)]), OrderedDict([('text', 'Regis Philbin, the program became a major ratings success throughout its initial summer run'), ('probability', 1.1565367338388186e-05), ('start_logit', 7.9423394203186035), ('end_logit', -3.7403860092163086)]), OrderedDict([('text', 'throughout its ABC tenure by Regis Philbin'), ('probability', 8.641061698026031e-06), ('start_logit', -3.7082467079162598), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Regis Philbin, the program became a major ratings success throughout its initial summer run,'), ('probability', 4.340281561528615e-06), ('start_logit', 7.9423394203186035), ('end_logit', -4.720461845397949)]), OrderedDict([('text', 'Millionaire, a game show based on the British program of the same title. Hosted throughout its ABC tenure by Regis Philbin'), ('probability', 2.581499164789754e-06), ('start_logit', -4.9164018630981445), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'game show based on the British program of the same title. Hosted throughout its ABC tenure by Regis Philbin'), ('probability', 1.5768661257791343e-06), ('start_logit', -5.409332752227783), ('end_logit', 7.618710517883301)]), OrderedDict([('text', 'Philbin,'), ('probability', 1.1642894253705209e-06), ('start_logit', -0.08077805489301682), ('end_logit', 1.9868273735046387)]), OrderedDict([('text', 'Hosted throughout its ABC tenure by Regis'), ('probability', 8.178505594859486e-07), ('start_logit', 2.464038133621216), ('end_logit', -0.9111754298210144)]), OrderedDict([('text', 'Meredith Vieira'), ('probability', 2.1176539810270548e-07), ('start_logit', 0.06787821650505066), ('end_logit', 0.13378390669822693)]), OrderedDict([('text', 'Hosted throughout its ABC tenure by Regis Phil'), ('probability', 1.7329348591843886e-07), ('start_logit', 2.464038133621216), ('end_logit', -2.4628684520721436)]), OrderedDict([('text', 'by Regis Philbin,'), ('probability', 1.1936571067088074e-07), ('start_logit', -2.358452320098877), ('end_logit', 1.9868273735046387)]), OrderedDict([('text', 'bin,'), ('probability', 1.0250649806025293e-07), ('start_logit', -2.5107181072235107), ('end_logit', 1.9868273735046387)])]\n",
      "-------------------------------------------------------------------------------------------------------------------\n",
      "Paragraph:\n",
      "In 2004, ABC's average viewership declined by ten ratings points, landing the network in fourth place, behind NBC, CBS and Fox (by the following year, the combined season-ending average audience share of ABC, NBC and CBS represented only 32% of U.S. households). However, during the 2004–05 season, the network experienced unexpected success with new series such as Desperate Housewives, Lost and Grey's Anatomy as well as reality series Dancing with the Stars, which helped ABC rise to second place, jumping ahead of CBS, but behind a surging Fox. On April 21, 2004, Disney announced a restructuring of its Disney Media Networks division with Anne Sweeney being named president of ABC parent Disney–ABC Television Group, and ESPN president George Bodenheimer becoming co-CEO of the division with Sweeney, as well as president of ABC Sports. On December 7, 2005, ABC Sports and ESPN signed an eight-year broadcast rights agreement with NASCAR, allowing ABC and ESPN to broadcast 17 Nextel Cup races each season (comprising just over half of the 36 races held annually) effective with the 2006 season.\n",
      "\n",
      "Question:\n",
      "Who was named president of Disney-ABC television group in 2004?\n",
      "\n",
      "Ground truth answers:\n",
      "['Anne Sweeney', 'Anne Sweeney', 'Anne Sweeney']\n",
      "\n",
      "Predicted answer:\n",
      "Anne Sweeney\n",
      "\n",
      "Top N best answers\n",
      "[OrderedDict([('text', 'Anne Sweeney'), ('probability', 0.9955952018782683), ('start_logit', 8.357232093811035), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Sweeney'), ('probability', 0.0036178995163272075), ('start_logit', 2.7397849559783936), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television Group,'), ('probability', 0.00015457112079484626), ('start_logit', 8.357232093811035), ('end_logit', -0.4140002131462097)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television Group'), ('probability', 0.00013177008073689449), ('start_logit', 8.357232093811035), ('end_logit', -0.5735959410667419)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television Group, and ESPN president George Bodenheimer becoming co-CEO of the division with Sweeney'), ('probability', 0.00010324790082242362), ('start_logit', 8.357232093811035), ('end_logit', -0.8175216317176819)]), OrderedDict([('text', 'Anne'), ('probability', 8.954744862928307e-05), ('start_logit', 8.357232093811035), ('end_logit', -0.9598858952522278)]), OrderedDict([('text', 'Disney announced a restructuring of its Disney Media Networks division with Anne Sweeney'), ('probability', 8.742559416104907e-05), ('start_logit', -0.9830758571624756), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'On April 21, 2004, Disney announced a restructuring of its Disney Media Networks division with Anne Sweeney'), ('probability', 5.3945594647410474e-05), ('start_logit', -1.4658879041671753), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television Group, and ESPN president George Bodenheimer'), ('probability', 5.264419046702813e-05), ('start_logit', 8.357232093811035), ('end_logit', -1.4910986423492432)]), OrderedDict([('text', 'April 21, 2004, Disney announced a restructuring of its Disney Media Networks division with Anne Sweeney'), ('probability', 2.4206634534182337e-05), ('start_logit', -2.2672371864318848), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Disney Media Networks division with Anne Sweeney'), ('probability', 2.0267835830057236e-05), ('start_logit', -2.444828748703003), ('end_logit', 8.356441497802734)]), OrderedDict([('text', '2004, Disney announced a restructuring of its Disney Media Networks division with Anne Sweeney'), ('probability', 1.651158022899612e-05), ('start_logit', -2.6498019695281982), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Anne Sweeney being named president'), ('probability', 1.5100346698474555e-05), ('start_logit', 8.357232093811035), ('end_logit', -2.7399368286132812)]), OrderedDict([('text', 'Anne Sweeney being named'), ('probability', 1.048714962404428e-05), ('start_logit', 8.357232093811035), ('end_logit', -3.104503870010376)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC'), ('probability', 7.763457777829983e-06), ('start_logit', 8.357232093811035), ('end_logit', -3.405226707458496)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC'), ('probability', 7.279500152601964e-06), ('start_logit', 8.357232093811035), ('end_logit', -3.469592332839966)]), OrderedDict([('text', 'with Anne Sweeney'), ('probability', 4.5077168678184705e-06), ('start_logit', -3.948073148727417), ('end_logit', 8.356441497802734)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television Group, and ESPN president George'), ('probability', 3.6106978117933016e-06), ('start_logit', 8.357232093811035), ('end_logit', -4.170753479003906)]), OrderedDict([('text', 'Anne Sweeney being named president of ABC parent Disney–ABC Television'), ('probability', 3.450058674894645e-06), ('start_logit', 8.357232093811035), ('end_logit', -4.216263294219971)]), OrderedDict([('text', 'Sweeney being named president of ABC parent Disney–ABC Television Group,'), ('probability', 5.616969448093071e-07), ('start_logit', 2.7397849559783936), ('end_logit', -0.4140002131462097)])]\n",
      "-------------------------------------------------------------------------------------------------------------------\n",
      "Paragraph:\n",
      "In addition, there are $2 million worth of other ancillary events, including a week-long event at the Santa Clara Convention Center, a beer, wine and food festival at Bellomy Field at Santa Clara University, and a pep rally. A professional fundraiser will aid in finding business sponsors and individual donors, but still may need the city council to help fund the event. Additional funding will be provided by the city council, which has announced plans to set aside seed funding for the event.\n",
      "\n",
      "Question:\n",
      "Where was a beer, wine and food festival held at prior to the Super Bowl?\n",
      "\n",
      "Ground truth answers:\n",
      "['Bellomy Field', 'Bellomy Field', 'Santa Clara Convention Center']\n",
      "\n",
      "Predicted answer:\n",
      "Bellomy Field at Santa Clara University\n",
      "\n",
      "Top N best answers\n",
      "[OrderedDict([('text', 'Bellomy Field at Santa Clara University'), ('probability', 0.8911798723290278), ('start_logit', 8.530954360961914), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Santa Clara University'), ('probability', 0.05653501302954448), ('start_logit', 5.773268222808838), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Bellomy Field'), ('probability', 0.04355919507587527), ('start_logit', 8.530954360961914), ('end_logit', 5.026169300079346)]), OrderedDict([('text', 'Bellomy Field at Santa Clara University,'), ('probability', 0.0041427350628246784), ('start_logit', 8.530954360961914), ('end_logit', 2.6734046936035156)]), OrderedDict([('text', 'at Bellomy Field at Santa Clara University'), ('probability', 0.0026144465466759878), ('start_logit', 2.699460506439209), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Bellomy Field at Santa Clara'), ('probability', 0.0007185308301180144), ('start_logit', 8.530954360961914), ('end_logit', 0.9215018153190613)]), OrderedDict([('text', 'Santa Clara University,'), ('probability', 0.00026280842737466214), ('start_logit', 5.773268222808838), ('end_logit', 2.6734046936035156)]), OrderedDict([('text', 'Santa Clara Convention Center, a beer, wine and food festival at Bellomy Field at Santa Clara University'), ('probability', 0.0002477774410439148), ('start_logit', 0.3431837260723114), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Field at Santa Clara University'), ('probability', 0.00016740039422735777), ('start_logit', -0.04895868897438049), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'at Bellomy Field'), ('probability', 0.0001277892271562228), ('start_logit', 2.699460506439209), ('end_logit', 5.026169300079346)]), OrderedDict([('text', 'at Santa Clara University'), ('probability', 0.00010075950695038713), ('start_logit', -0.556610643863678), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Bellomy Field at'), ('probability', 5.6814224825865046e-05), ('start_logit', 8.530954360961914), ('end_logit', -1.615920066833496)]), OrderedDict([('text', 'Bellomy Field at Santa Clara University, and a pep rally.'), ('probability', 5.503602065575232e-05), ('start_logit', 8.530954360961914), ('end_logit', -1.647718906402588)]), OrderedDict([('text', 'University'), ('probability', 4.749428214965224e-05), ('start_logit', -1.308737874031067), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Santa Clara'), ('probability', 4.55824363904097e-05), ('start_logit', 5.773268222808838), ('end_logit', 0.9215018153190613)]), OrderedDict([('text', 'Bellomy'), ('probability', 3.587220941206712e-05), ('start_logit', 8.530954360961914), ('end_logit', -2.0757439136505127)]), OrderedDict([('text', 'Bellomy Field at Santa'), ('probability', 3.2615110636404645e-05), ('start_logit', 8.530954360961914), ('end_logit', -2.170931100845337)]), OrderedDict([('text', 'beer, wine and food festival at Bellomy Field at Santa Clara University'), ('probability', 2.7231470259308925e-05), ('start_logit', -1.8649739027023315), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Clara University'), ('probability', 2.310549787851485e-05), ('start_logit', -2.0292766094207764), ('end_logit', 8.044594764709473)]), OrderedDict([('text', 'Bell'), ('probability', 1.992087697312563e-05), ('start_logit', 8.530954360961914), ('end_logit', -2.663938522338867)])]\n",
      "-------------------------------------------------------------------------------------------------------------------\n"
     ]
    }
   ],
   "source": [
    "for i in [0, 10, 100]:\n",
    "    print('Paragraph:')\n",
    "    print(dev_df.iloc[i]['doc_text'])\n",
    "    print()\n",
    "    print('Question:')\n",
    "    print(dev_df.iloc[i]['question_text'])\n",
    "    print()\n",
    "    print('Ground truth answers:')\n",
    "    print(dev_df.iloc[i]['answer_text'])\n",
    "    print()\n",
    "    print('Predicted answer:')\n",
    "    print(final_answers[dev_df.iloc[i]['qa_id']])\n",
    "    print()\n",
    "    print('Top N best answers')\n",
    "    print(nbest_answers[dev_df.iloc[i]['qa_id']])\n",
    "    print('-------------------------------------------------------------------------------------------------------------------')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluate"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Question answering task is usually evaluated on two metrics: exact match (EM) and F1 score.   \n",
    "The exact match is computed by first performing some simple normalization (e.g. remove punctuation and convert to lower case) on the ground truth and predicted answers and check if they match exactly after normalization.   \n",
    "F1 score is computed from token-level precision and recall by comparing the ground truth and predicted answers. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"exact\": 86.6414380321665,\n",
      "  \"f1\": 92.68221713064221,\n",
      "  \"total\": 10570,\n",
      "  \"HasAns_exact\": 86.6414380321665,\n",
      "  \"HasAns_f1\": 92.68221713064221,\n",
      "  \"HasAns_total\": 10570\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "evaluation_result = evaluate_qa(actual_dataset=dev_dataset,\n",
    "                                preds=final_answers)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The table below compares running time and model performance of BERT, XLNet, and DistilBert on Standard_NC24rs_v3 DSVM.\n",
    "\n",
    "|Model name|Training time|Scoring time|Exact Match (EM)|F1 score|\n",
    "|:---------|:------------|:-----------|:--------------|--------|\n",
    "|bert-large-cased-whole-word-masking| 3.4 hrs| ~ 5 mins|86.64|92.68|\n",
    "|xlnet-large-cased|5.2 hrs| ~ 10 mins|84.67|91.69|\n",
    "|distilbert-base-uncased|0.66 hr| ~ 1 min|76.62|84.71|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sb.glue(\"exact\", evaluation_result[\"exact\"])\n",
    "sb.glue(\"f1\", evaluation_result[\"f1\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## References\n",
    "\n",
    "1. Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, Percy Liang, [*SQuAD: 100,000+ Questions for Machine Comprehension of Text*](https://arxiv.org/abs/1606.05250), EMNLP, 2016.\n",
    "2. Pranav Rajpurkar, Robin Jia, Percy Liang, [*Know What You Don't Know: Unanswerable Questions for SQuAD*](https://arxiv.org/abs/1806.03822), ACL, 2018"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "nlp_gpu",
   "language": "python",
   "name": "nlp_gpu"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
